import cv2
import numpy as np
import os
from tqdm import tqdm
import json
import shutil
from utils.ddd_utils import draw_box_3d, project_to_image, compute_box_3d, rot_y2alpha, alpha2rot_y
from pathlib import Path
from utils.yolo_test import Yolodect

print("version", 1)
# kitti ppt
# class dz dy dx x y z r
# show_image = True
show_image = True


def mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)


def loadlabel(label_path):
    with open(label_path) as f:
        labels = json.load(f)
        objs = []
        for label in labels:
            obj = {}
            position = label['psr']['position']
            rotation = label['psr']['rotation']['z']
            scale = label['psr']['scale']
            x, y, z = position["x"], position["y"], position["z"]
            xs, ys, zs = scale["x"], scale["y"], scale["z"]
            obj["obj_type"], obj["obj_id"] = label['obj_type'], label['obj_id']
            obj["box"] = [x, y, z, xs, ys, zs, rotation]
            objs.append(obj)
    return objs


def save_label(labels_list, label_path):
    for labels in labels_list:
        path_split = label_path.split("/")
        file_name = "{}_c25/training/label/{}_{}.txt".format(path_split[0], path_split[3].split(".")[0], labels[1])
        with open(file_name, "w") as f:
            if len(labels[0]) == 0:
                s = "others 0 0 0 0 0 0 0 0 0 0 0 0 0 0.0"
                f.write(s)
            for label in labels[0]:
                s = ""
                for l in label:
                    s += str(l)
                    if not l is label[-1]:
                        s += " "
                if not label is labels[0][-1]:
                    s += "\n"
                f.write(s)


def save_calib(calibs, label_path):
    for cam in calibs:
        calib = calibs[cam]
        s = "P0:"
        for c in calib["intrinsic"]:
            s += " "
            s += str(c)
        s1 = "P1:"
        for d in calib["extrinsic"]:
            s1 += " "
            s1 += str(d)

        path_split = label_path.split("/")
        file_name = "{}_c25/training/calib/{}_{}.txt".format(path_split[0], path_split[3].split(".")[0], cam)
        with open(file_name, "w") as f:
            f.write(s)
            f.write('\r\n')
            f.write(s1)


def save_image(label_path):
    path_split = label_path.split("/")
    for cam in ["right"]:
        source_image_path = label_path.replace("/label", "/camera/{}".format(cam)).replace(".json", ".png")
        target_image_path = "{}_c25/training/images/{}_{}.png".format(path_split[0], path_split[3].split(".")[0], cam)
        shutil.copy(source_image_path, target_image_path)


def readcalib(bag_dir):
    calibs = {}
    for cam in ["right"]:
        with open(os.path.join(bag_dir, "calib/camera/" + cam + ".json")) as f:
            calib = json.load(f)
        # calib_ex = calib["extrinsic"]
        # calib_in = calib["intrinsic"]
        calib["intrinsic"].insert(4, 0.0)
        calib["intrinsic"].insert(8, 0.0)
        calib["intrinsic"].insert(12, 0.0)
        calibs[cam] = calib
    return calibs


def init_dir(dataname):
    mkdir("./{}_c25/".format(dataname))
    mkdir("./{}_c25/training/".format(dataname))
    mkdir("./{}_c25/training/images/".format(dataname))
    mkdir("./{}_c25/training/images/trainval".format(dataname))
    mkdir("./{}_c25/training/label/".format(dataname))
    mkdir("./{}_c25/training/calib/".format(dataname))


def get_imgs_path(label_path):
    imgs_path = {}
    for cam in ["right"]:
        img_path = label_path.replace("label", "camera/{}".format(cam)).replace(".json", ".png")
        imgs_path[cam] = img_path
    return imgs_path


# def cpimgs(dataname,bag_name,imgs_path):
#     for cam in ["right","right","right"]:
#         img_path =imgs_path[cam]
#         img_name = os.path.basename(img_path)
#         target_img_name = "{}_{}".format(bag_name,img_name).replace(".png","_{}.png".format(cam))
#         target_img_path = "./{}_c25/training/trainval/images/{}".format(dataname,target_img_name)
#         shutil.copy(img_path, target_img_path)

def compare_bbox(b2s, b3):
    min_loss = 1000
    limit_error = 300
    b3_c, b3_xy0, b3_xy1 = b3.mean(0), [max(b3[:, 0]), max(b3[:, 1])], [min(b3[:, 0]), min(b3[:, 1])]

    # h = max(b3[:, 1]) - min(b3[:, 1])
    # w = max(b3[:, 0]) - min(b3[:, 0])
    # for b2 in b2s:
    #     b2c = np.array([b2[0][2]+b2[0][0],b2[0][3]+b2[0][1]])/2
    #     b2_xy0 = np.array([b2[0][0],b2[0][1]])
    #     b2_xy1 = np.array(([b2[0][2],b2[0][3]]))
    #     # b2_c = (b2_xy0+b2_xy1)/2
    #     # w,h = b2_xy1[0] - b2_xy0[0],b2_xy1[1] - b2_xy0[1]
    #     # w3,h3 = b3_xy1[0] - b3_xy0[0],b3_xy1[1] - b3_xy0[1]
    #     loss = (abs((b3_xy0-b2_xy0)).mean()+abs((b3_xy1-b2_xy1)).mean())+np.square(b3_c-b2c).mean()
    #     # print(loss,w,h)
    #     if loss <min_loss and loss<limit_error:
    #         min_loss = loss
    #         bbox_2d =[b2_xy0,b2_xy1]
    # # print(min_loss)
    # if min_loss<limit_error:
    #     return bbox_2d
    return [b3_xy0, b3_xy1]


# names = ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'rider', 'others']
# class_map = {"Car":"car","Pedestrian":"person","Truck":"truck",
#              "Bus":"bus","Motorcycle":"motorcycle","MotorcyleRider":"rider",
#              "Bicycle":"bicycle","BicycleRider":"rider","Van":"truck","Unknown":"person"} #,"Van":"van",

class_map = {"Car": "Car", "Pedestrian": "Pedestrian", "Truck": "Truck",
             "Bus": "Bus", "Motorcycle": "Motorcyclist", "MotorcyleRider": "Motorcyclist",
             "Bicycle": "Cyclist", "BicycleRider": "Cyclist", "Van": "Van", "Unknown": "Trafficcone"}  # ,"Van":"van",


# "Car", "Truck", "Van", "Bus", "Pedestrian", "Cyclist", "Tricyclist", "Motorcyclist", "Barrowlist",
#             "Trafficcone"
def get_label(boxs_2d, boxs_3d):
    # bboxs_3d.append([dim, location, rotation_y, box_2d[-1],class_name,obj_id])
    # col1 ： 类标签
    # col5-8 ：2D bbox，左上和右下坐标
    # col9-11 : 3D bbox dim 高、宽、长
    # col12-14 ：3D bbox centor 坐标
    # col15: 车体方向角
    # print(boxs_2d)
    labels = []
    for box_3d in boxs_3d:
        class_name = class_map[box_3d[4]]
        # obj_id = box_3d[5]
        bbox_from3d = box_3d[3]
        [bbox_xy0, bbox_xy1] = compare_bbox(boxs_2d, bbox_from3d)
        dim = box_3d[0]
        location = [_[0] for _ in box_3d[1].tolist()]
        rotation = box_3d[2]
        label = [class_name, 0, 0, 0, bbox_xy0[0], bbox_xy0[1], bbox_xy1[0], bbox_xy1[1]] + dim + location + [rotation]
        labels.append(label)
    return labels


out_size = 400
world_size = 80


def project_3d_to_bird(pt):
    pt[0] += world_size / 2
    pt[1] = world_size - pt[1]
    pt = pt * out_size / world_size
    return pt.astype(np.int32)


def showddd(imgs_path, calibs, objs):
    imgs = None
    bird_views = None
    labels_list = []

    while (1):
        quit = 0
        for i in range(-10, 10):

            flag = 0
            for cam in ["right"]:
                # for cam in ["right"]:
                bboxs_3d = []

                # str((int(imgs_path[cam].split("/")[-1].split('.')[0]) + i)).zfill(6)
                cur_img_path = os.path.join(imgs_path['right'])
                img = cv2.imread(cur_img_path)

                if img is None:
                    continue

                bird_view = np.ones((out_size, out_size, 3), dtype=np.uint8) * 230

                for obj in objs:
                    [x, y, z, xs, ys, zs, r] = obj["box"]
                    z -= zs / 2
                    rotation_y = -r
                    if cam == "right":
                        rotation_y -= 2.35  # 不同视角需要修改这个
                    dim = [zs, ys, xs]
                    class_name = obj["obj_type"]
                    # obj_id = obj["obj_id"]


                    location = [x, y, z]
                    location = np.expand_dims(np.concatenate([location, np.ones(1)], 0), 1)
                    p = np.array(calibs[cam]["extrinsic"]).reshape((4, 4))[:3]
                    location = np.dot(p, location)
                    # print(class_name, location.T)

                    calib = np.array(calibs[cam]["intrinsic"]).reshape((3, 4))
                    box_3d = compute_box_3d(dim, location, rotation_y)
                    rect = box_3d[:4, [0, 2]]
                    for k in range(4):
                        rect[k] = project_3d_to_bird(rect[k])
                    lc = (250, 152, 12)
                    cv2.polylines(
                        bird_view, [rect.reshape(-1, 1, 2).astype(np.int32)],
                        True, lc, 2, lineType=cv2.LINE_AA)

                    if box_3d.mean(0)[2] < 0:
                        continue
                    box_2d = project_to_image(box_3d, calib)
                    h = max(box_2d[:, 1]) - min(box_2d[:, 1])
                    w = max(box_2d[:, 0]) - min(box_2d[:, 0])
                    if abs(h) < 10 or abs(w) < 10:
                        continue
                        # print("wh",class_name,box_2d[0][0]-box_2d[7][0],box_2d[0][1]-box_2d[7][1])
                    if min(box_2d[-1]) < -100 or box_2d[-1][0] > 1400 or box_2d[-1][1] > 900:
                        continue
                    img = draw_box_3d(img, box_2d)

                    bboxs_3d.append([dim, location, rotation_y, box_2d, class_name])

                bbox_2d = []
                labels = get_label(bbox_2d, bboxs_3d)

                for l in labels:
                    bbox = l[4:8]
                    xyxy = bbox
                    # print((xyxy[0],xyxy[1]))
                    cv2.rectangle(img, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (255, 255, 0), 2)

            if show_image:
                if img is None:
                    continue

                cv2.imshow(imgs_path[cam].split("/")[-1] + "----" + cur_img_path.split('/')[-1], img)
                keyValue = cv2.waitKey()
                if keyValue & 0xFF == ord('q'):  # ‘q’键实现退出
                    exit(0)
                elif keyValue & 0xFF == ord('s'):
                    # cv2.imwrite(imgs_path[cam].split("/")[-1] + "----" + cur_img_path.split('/')[-1], img)
                    path_split = cur_img_path.split('/')
                    target_image_path = "{}_c25/training/images/{}_{}.png".format(path_split[0],

                                                                                  path_split[-1].split(".")[
                                                                                      0],
                                                                                  cam)
                    cv2.imwrite(target_image_path, img)
                    labels_list.append([labels, cam])
                    flag = 1
                    quit = 1
                    cv2.destroyAllWindows()

                elif keyValue & 0xFF == ord('s'):
                    i -= 1
                    cv2.destroyAllWindows()
                elif keyValue & 0xFF == ord('n'):
                    quit = 1
                    cv2.destroyAllWindows()
                else:
                    cv2.destroyAllWindows()

            if flag or quit:
                break

        if quit:
            break

    return labels_list


def main():
    dataname = "8016"
    bags_name = os.listdir(dataname)
    init_dir(dataname)
    # yolodete = Yolodect()

    for bag_idx, bag_name in enumerate(bags_name):
        bag_dir = os.path.join(dataname, bag_name)
        lidar_dir = "{}/lidar/".format(bag_dir)
        label_dir = "{}/label/".format(bag_dir)
        lidar_files = os.listdir(lidar_dir)
        label_files = os.listdir(label_dir)

        # lidar_files.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
        bar = tqdm(lidar_files)

        label_num = len(label_files)
        lidar_num = len(lidar_files)

        calibs = readcalib(bag_dir)

        for lidar_file in bar:
            bar.set_description("bag_name bag/all {} {}/{}".format(bag_name, bag_idx + 1, len(bags_name)))
            if label_num < lidar_num:
                print("label_num  lidar_num", label_num, lidar_num)
                break
            label_file = lidar_file.replace(".bin", ".json")
            label_path = os.path.join(label_dir, label_file)
            # print(label_path)

            objs = loadlabel(label_path)

            imgs_path = get_imgs_path(label_path)
            # cpimgs(dataname,bag_name,imgs_path)

            labels_list = showddd(imgs_path, calibs, objs)
            save_label(labels_list, label_path)
            save_calib(calibs, label_path)
            # save_image(label_path)


if __name__ == '__main__':
    main()
